import json
import uuid
import time
import redis
import yaml

CONFIG_FILE = "./common/config.yaml"
def load_config():
    """
    从YAML配置文件加载配置

    Returns:
        dict: 配置字典
    """
    try:
        with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        return config
    except Exception as e:
        import sys
        sys.exit(1)

config = load_config()

redis_client = redis.Redis(
    host=config["redis"]["host"],
    port=config["redis"]["port"],
    password=config["redis"]["password"],
    db=config["redis"]["db"],
    decode_responses=True,
)

def get_redis():
    return redis_client

def set_proxies(proxies):
    for proxy in proxies:
        redis_client.sadd('bilibili_scrawler:proxies', json.dumps(proxy))
    redis_client.expire('bilibili_scrawler:proxies', 270)

def get_proxies():
    proxies = []
    for item in redis_client.smembers('bilibili_scrawler:proxies'):
        proxies.append(json.loads(item))
    return proxies

def get_lock(lock_key, lock_timeout=10):
    identifier = uuid.uuid4()
    if redis_client.set(lock_key, str(identifier), nx=True, ex=lock_timeout):
        return True, str(identifier)
    return False,""

def release_lock(lock_key: str, identifier: str) -> bool:
    lua_script = """
    if redis.call("GET", KEYS[1]) == ARGV[1] then
        return redis.call("DEL", KEYS[1])
    else
        return 0
    end
    """
    return bool(redis_client.eval(lua_script, 1, lock_key, identifier))

def delete_lock(lock_key):
    redis_client.delete(lock_key)

def get_trino_lock(lock_key, lock_timeout=10):
    identifier = uuid.uuid4()
    while True:
        if redis_client.set(lock_key, str(identifier), nx=True, ex=lock_timeout):
            print("获取trino锁成功")
            return True, str(identifier)
        print("获取trino锁失败，正在重试...")
        time.sleep(1)



